import torch
import torch.nn as nn
from torchvision.transforms import ToTensor, Compose, Resize
from PIL import Image, ImageTk
import torch.nn.functional as F
import tkinter as tk
import pyautogui

def load_trained_model(model_path='mnist_model.pth', num_classes=10):# 加载训练好的模型并返回模型对象
    class Model(nn.Module):
        # 自定义的神经网络模型类,继承自PyTorch的nn.Module基类,用于定义模型的结构和前向传播逻辑。
        def __init__(self, num_classes):
            super().__init__()
            # 特征提取网络
            self.conv1 = nn.Conv2d(1, 32, kernel_size=3)# 定义第一个卷积层
            self.pool1 = nn.MaxPool2d(2)# 定义第一个最大池化层
            self.conv2 = nn.Conv2d(32, 64, kernel_size=3)# 定义第二个卷积层
            self.pool2 = nn.MaxPool2d(2)# 定义第二个最大池化层
            # 分类网络
            self.fc1 = nn.Linear(1600, 64)# 定义第一个全连接层
            self.fc2 = nn.Linear(64, num_classes)# 定义第二个全连接层

        def forward(self, x):# 定义前向传播逻辑
            x = self.pool1(F.relu(self.conv1(x)))# 第一个卷积层+ReLU+池化
            x = self.pool2(F.relu(self.conv2(x)))# 第二个卷积层+ReLU+池化
            x = torch.flatten(x, start_dim=1)# 将特征图展平成一维向量
            x = F.relu(self.fc1(x))# 第一个全连接层+ReLU
            x = self.fc2(x)# 第二个全连接层
            return x

    model = Model(num_classes).to(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))# 定义模型对象并加载到GPU或CPU
    model.load_state_dict(torch.load(model_path))# 加载模型参数
    model.eval()  # 将模型设置为评估模式,用于后续预测
    return model


def recognize_digit_in_image(image_path='input_digit_image.png'):# 识别图像中的数字并返回数字值
    # 加载模型
    trained_model = load_trained_model()
    # 定义图像预处理的transforms
    transform = Compose([
        Resize((28, 28)),
        ToTensor(),
    ])
    image = Image.open(image_path).convert('L')
    input_tensor = transform(image).unsqueeze(0).to(torch.device('cuda' if torch.cuda.is_available() else 'cpu'))

    # 使用模型进行预测
    with torch.no_grad():
        output = trained_model(input_tensor)
        prediction = output.argmax(1).item()
        print(f"预测结果为：{prediction}")
        return prediction


class HandwrittenDigitRecognizer(tk.Tk):
    def __init__(self):
        super().__init__()
        self.title("手写数字识别")
        self.geometry("760x400+1200+700")
        # 禁止调整窗口大小
        self.resizable(0, 0)
        # 创建左侧画布用于手写数字
        self.left_canvas = tk.Canvas(self, width=400, height=400, bg='black')
        self.left_canvas.pack(side=tk.LEFT)
        # 创建右侧主框架,用于承载按钮框架和标签框架
        self.right_frame = tk.Frame(self)
        self.right_frame.pack(side=tk.RIGHT, fill=tk.BOTH, expand=True)
        # 创建上部框架,用于放置两个按钮
        self.button_frame = tk.Frame(self.right_frame,bg='white')
        self.button_frame.pack(side=tk.TOP, fill=tk.X)
        # 创建确认按钮
        self.confirm_button = tk.Button(self.button_frame, text="确认", command=self.predict_digit, font=("黑体", 16),relief="groove",bg='white')
        self.confirm_button.pack(side=tk.LEFT, padx=50,pady=10)
        # 创建清屏按钮
        self.clear_button = tk.Button(self.button_frame, text="清屏", command=self.clear_canvas, font=("黑体", 16),relief="groove",bg='white')
        self.clear_button.pack(side=tk.LEFT, padx=30,pady=10)
        # 创建下部框架,用于放置显示预测结果的Label
        self.label_frame = tk.Frame(self.right_frame,bg='white')
        self.label_frame.pack(side=tk.BOTTOM, fill=tk.X)
        # 用于记录显示图片的Label组件,初始化为None
        self.display_panel = tk.Label(self.label_frame, text="",width=250, height=250, bg='white')
        self.display_panel.pack(side="top", expand="yes")
        # 用于显示预测结果的Label
        self.result_label = tk.Label(self.label_frame, text="", font=("黑体", 16),bg='white')
        self.result_label.pack(side="bottom",pady=10)
        # 用于记录鼠标绘制的点
        self.points = []
        self.left_canvas.bind("<B1-Motion>", self.draw)
        self.left_canvas.bind("<ButtonRelease-1>", self.clear_point)  # 绑定鼠标按键抬起事件
        # 判断是否绘画
        self.drawing = False
        
    def draw(self, event):#在左侧画布上绘制线条,模拟手写数字
        x, y = event.x, event.y
        self.points.append((x, y))
        if len(self.points) > 1:
            self.drawing = True
            x1, y1 = self.points[-2]
            self.left_canvas.create_oval(x1, y1, x, y, width=45,fill='white', outline='white')

    def predict_digit(self):# 按下确认开始识别
        if not self.drawing:
            print("请先绘制数字")
            return
        # 禁用确认按钮和清屏按钮,防止重复操作
        self.confirm_button.config(state=tk.DISABLED)
        self.clear_button.config(state=tk.DISABLED)
        # 截图获取画布内容并保存为PNG格式
        x = self.left_canvas.winfo_rootx()
        y = self.left_canvas.winfo_rooty()
        width = self.left_canvas.winfo_width()
        height = self.left_canvas.winfo_height()
        screenshot = pyautogui.screenshot(region=(x+2, y+2, width-5, height-5))
        # 处理图像
        img = screenshot.convert('L')# 转换为灰度图
        img.save('input_digit_image.png')
        img = img.resize((260, 260))  
        img = img.point(lambda x: 255 - x if x < 255 else 0)# 反转黑白
        img_tk = ImageTk.PhotoImage(img)# 转换为Tkinter可用的PhotoImage对象
        # 在tkinter中显示图片
        if self.display_panel:
            self.display_panel.destroy()
        self.display_panel = tk.Label(self.label_frame, image=img_tk,bd=0)
        self.display_panel.image = img_tk  # 保持对图像对象的引用,防止被垃圾回收
        self.display_panel.pack(side="top", expand="yes")
        # 删除临时文件
        prediction = recognize_digit_in_image()
        # 显示预测结果并自动清空画布
        self.result_label.config(text=f"预测结果: {prediction}")
        self.clear_canvas()
        # 启用确认按钮和清屏按钮,允许继续操作
        self.confirm_button.config(state=tk.NORMAL)
        self.clear_button.config(state=tk.NORMAL)

    def clear_canvas(self):# 清空左侧画布上绘制的内容以及记录绘制点的列表
        self.left_canvas.delete("all")
        self.points = []
        self.drawing = False

    def clear_point(self, event):  # 在鼠标按键抬起时（一笔绘画结束）,清除记录的坐标点
        self.points=[] 

if __name__ == "__main__":
    print("----- "*7+"-----")
    print("|       ☆ 欢迎使用手写数字识别器! ☆           |")
    print("|     *****本程序基于PyTorch实现*****         |")
    print("|  使用MNIST数据集训练的模型进行手写数字识别  |")
    print("|                 作者:Zaly                   |")
    print("-----+"*7+"-----")
    app = HandwrittenDigitRecognizer()
    app.mainloop()